from typing import Annotated

from langchain_core.tools import tool
from langgraph.constants import START, END
from langgraph.graph import add_messages, StateGraph
from typing_extensions import TypedDict

from src.ai.langchain.init_llm import get_llm


llm = get_llm()


class State(TypedDict):
    messages: Annotated[list, add_messages]


def chatbot(state: State):
    return {"messages": [llm.invoke(state['messages'])]}


graph_builder = StateGraph(State)


graph_builder.add_node("chatbot", chatbot)
graph_builder.add_edge(START, "chatbot")
graph_builder.add_edge("chatbot", END)
graph = graph_builder.compile()

# 生成图
# png = graph.get_graph().draw_mermaid_png()
#
# with open("langgraph_chatbot.png", "wb") as f:
#     f.write(png)

history = []


def stream_graph_updates(user_input: str, history: list):

    history += [{"role": "human", "content": user_input}]

    print("完成的请求消息：", history)

    result = graph.invoke({"messages": history})

    print("AI原始结果：", result)

    result_content = result["messages"][-1].content

    print("智能助手：", result_content)

    history.append({"role": "assistant", "content": result_content})

    # for event in graph.stream({"messages": [{"role": "human", "content": user_input}]}):
    #     for value in event.values():
    #         # print("聊天助手:", value["messages"][-1].content)
    #         print("聊天助手:", value["messages"])


while True:
    try:
        user_input = input("请输入: ")
        if user_input.lower() in ["quit", "exit", "q"]:
            print("Goodbye!")
            break
        stream_graph_updates(user_input, history)
    except:
        # fallback if input() is not available
        user_input = "What do you know about LangGraph?"
        print("异常，用户输入内容为: " + user_input)
        stream_graph_updates(user_input, history)
        break


